'''
PET 图像HU值转换
PET 图像标准化
PETSlice 生成
补充PETSlice文件地址到csv文件

裁剪PET,并resize至512x512大小，使之和CT尺寸保持一致，方便运用坐标计算suv值
'''


import pandas
import pydicom
import cv2
import numpy as np
import csv
import os


# 为便于之后的suv值计算，并没有进行归一化
#经统计，pet数据都分布在0 100000之间
def normalize_hu(image):
    MIN_BOUND = 0
    MAX_BOUND = 100000
    image = (image - MIN_BOUND) / (MAX_BOUND - MIN_BOUND)
    image[image > 1] = 1.
    image[image < 0] = 0.
    return image



def save_origin_pet():
    origin_path = 'D:/lung_cancer/data/all_data.csv'
    data = pandas.read_csv(origin_path)
    for i in range(len(data)):
        patient = data['patientID'][i]
        pet_path = data['PET_origin_path'][i]
        slice = pydicom.read_file('H:/'+pet_path)
        pet_array = slice.pixel_array


        save_path = 'D:/lung_cancer/data/Slice/origin_PETSlice/'+str(patient)+'_'+'PETSlice.npy'
        np.save(save_path, pet_array)
        print('%d--->patient: %s' % (i+1, str(patient)))

def save_cut_PET():
    origin_path = 'D:/lung_cancer/data/data_augmentation/divide_csv/over_sampling_five/augement_test2.csv'
    data = pandas.read_csv(origin_path)
    for i in range(len(data)):
        patient = data['patientID'][i]
        instancenumber = str(data['z'][i])
        ct_size = data['ct_size'][i]
        pet_size = data['pet_size'][i]
        ct_spacing = data['ct_x_spacing'][i]
        pet_spacing = data['pet_x_spacing'][i]
        pet_name = str(data['PET_origin_path'][i]).split('/')[-1]
        cancer_type = data['cancer_type'][i]

        pet_path = 'D:/lung_cancer/data/data_augmentation/origin_data_circle/'+str(patient)+'/PET/'+pet_name


        slice = pydicom.read_file(pet_path)
        pet_array = slice.pixel_array





        # print(pet_array.shape)

        real_size = int(round(pet_size*pet_spacing/ct_spacing))
        border = (real_size-ct_size)//2

        # print(real_size)
        # print(border)

        resized_array = cv2.resize(pet_array, (real_size, real_size))
        new_pet = resized_array[border:border+ct_size, border:border+ct_size]




        save_path = 'D:/lung_cancer/data/data_augmentation/divide_csv/over_sampling_five/Slice/'+str(patient)+'/PETSlice/'
        if not os.path.exists(save_path):
            os.makedirs(save_path)
        np.save(save_path+instancenumber+'.npy', new_pet)
        print('%d--->patient: %s' % (i+1, str(patient)))

# 将生成的pet_array补充到all_data1.csv中, 生成all_data2.csv
def add_petpath():
    data_path = 'D:/lung_cancer/data/data_augmentation/divide_csv/over_sampling_five/augement_test2.csv'
    f = csv.reader(open(data_path, 'r'))
    data = []
    for i in f:
        data.append(i)

    new_data = []

    for line in data[1:]:
        patient = str(line[0])
        cancer_type = str(line[5])
        instancenumber = str(line[1])
        pet_path = 'Slice/'+patient+'/PETSlice/'+instancenumber+'.npy'
        # origin_pet_path = 'Slice/origin_PETSlice/' + patient + '_PETSlice.npy'
        line.append(pet_path)
        # line.append(origin_pet_path)
        new_data.append(line)
    data[0].append('PETSlice_Path')
    # data[0].append('origin_PETSlice_Path')
    df = pandas.DataFrame(new_data, columns=data[0])
    df.to_csv('D:/lung_cancer/data/data_augmentation/divide_csv/over_sampling_five/augement_test3.csv', index=False)



if __name__ == '__main__':
    # save_cut_PET()
    add_petpath()